/*
 * Copyright (c) 2020-2099 the original author or authors. All rights reserve.
 */
package org.tianyun.cloud.datasource.sharding;

import org.apache.shardingsphere.core.yaml.swapper.ShardingRuleConfigurationYamlSwapper;
import org.apache.shardingsphere.shardingjdbc.api.ShardingDataSourceFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.tianyun.cloud.datasource.WrappedDataSource;
import org.tianyun.cloud.datasource.configuration.HikariProperties;
import org.tianyun.cloud.datasource.configuration.ShardingDataSourceProperties;
import org.tianyun.cloud.datasource.hikari.HikariDataSourceBuilder;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

/**
 * ShardingDataSourceBuilder
 *
 * @auther ebert_chan
 */
public class ShardingDataSourceBuilder {

    private static final Logger LOGGER = LoggerFactory.getLogger(ShardingDataSourceBuilder.class);

    private ShardingDataSourceProperties sharding;

    private boolean wrapped;

    private ShardingDataSourceBuilder(ShardingDataSourceProperties sharding) {
        this.sharding = sharding;
    }

    public static ShardingDataSourceBuilder configure(ShardingDataSourceProperties sharding) {
        Assert.notEmpty(sharding.getNodes(), "Failed to create the datasource, this argument 'datasource.sharding.nodes' is required; it must not be empty");

        LOGGER.info("----------------------------------------------------------------------------------------------");
        for (Map.Entry<String, HikariProperties> entry : sharding.getNodes().entrySet()) {
            String jdbcUrl = entry.getValue().getJdbcUrl();
            LOGGER.info("Sharding Datasource nodes[{0}]: {1}", entry.getKey(), StringUtils.hasText(jdbcUrl) ? jdbcUrl.substring(0, jdbcUrl.indexOf("?")) : "");
        }
        LOGGER.info("----------------------------------------------------------------------------------------------");

        ShardingDataSourceBuilder builder = new ShardingDataSourceBuilder(sharding);

        return builder;
    }

    public ShardingDataSourceBuilder wrapped() {
        this.wrapped = true;
        return this;
    }

    public DataSource build() throws SQLException {
        HikariProperties hikariProperties = null;
        Map<String, HikariProperties> nodes = sharding.getNodes();
        Map<String, DataSource> dataSourceMap = new HashMap<>(nodes.size());

        Boolean hasError = Boolean.FALSE;
        try {
            for (Map.Entry<String, HikariProperties> entry : nodes.entrySet()) {
                hikariProperties = entry.getValue();
                DataSource datasource = HikariDataSourceBuilder.configure(hikariProperties, sharding.getProps()).build();
                dataSourceMap.put(entry.getKey(), datasource);
            }
        } catch (Exception e) {
            hasError = Boolean.TRUE;
            throw e;
        } finally {
            if (hasError) {
                for (Map.Entry<String, DataSource> entry : dataSourceMap.entrySet()) {
                    try {
                        closeDataSource(entry.getValue());
                    } catch (Exception e) {
                        // ignore e
                    }
                }
            }
        }

        DataSource dataSource = ShardingDataSourceFactory.createDataSource(dataSourceMap, new ShardingRuleConfigurationYamlSwapper().swap(sharding.getRule()), sharding.getProps());
        if (!wrapped) {
            return dataSource;
        }

        WrappedDataSource wds = new WrappedDataSource(dataSource);
        wds.setDriverClassName(hikariProperties.getDriverClassName());
        wds.setJdbcUrl(hikariProperties.getJdbcUrl());

        return wds;
    }

    private void closeDataSource(DataSource dataSource) throws IOException {
        if (dataSource == null || !(dataSource instanceof Closeable)) {
            return;
        }
        Closeable closeable = (Closeable) dataSource;
        closeable.close();
    }

}
